import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class Learner {
	
	public static TrainingResult train(DocumentsLoader loader) {
		List<Document> documents = loader.getTrainingDocuments();
		int[] classes = loader.getClasses();
		
		Set<String> vocabulary = extractVocabulary(documents);
		int n = documents.size();	
		
		int[] prior = new int[classes.length];
		Map<String, double[]> condprob = new HashMap<String, double[]>();			
		
		for (int clazz : classes) {
			
			Map<String, Integer> tokensCount = new HashMap<String, Integer>();
			int nc = loader.trainingDocumentsCount(clazz);
			prior[clazz] = nc / n;
			
			for (String t : vocabulary) {
				tokensCount.put(t, countTokens(loader, clazz, t));				
			}
			
			for (String t : vocabulary) {
				double sum = 0;
				for (String tfree : vocabulary) {
					sum += tokensCount.get(tfree) + 1;
				}
				
				double val = (double) (tokensCount.get(t) + 1) / (double) sum;
				
				if (condprob.containsKey(t)) {
					double[] c = condprob.get(t);
					c[clazz] = val;
				} else {
					double[] c = new double[classes.length];
					c[clazz] = val;
					condprob.put(t, c);
				}
			}
		}
		return new TrainingResult(vocabulary, prior, condprob);
	}
	
	public static int apply(int[] classes, TrainingResult result, Document d) {
		Set<String> vocab = result.vocabulary;
		List<String> tokens = d.getTokens();
		List<String> words = new ArrayList<String>();
		
		for(String t : tokens) {
			if (vocab.contains(t)) {
				words.add(t);
			}
		}
		
		double[] score = new double[classes.length];
		
		int[] prior = result.prior;
		Map<String, double[]> condprob = result.condprob;
		
		for (int c : classes) {
			score[c] = Math.log(prior[c]);
			for (String t : words) {
				score[c] += Math.log(condprob.get(t)[c]);
			}
		}
		
		int maxi = 0;
		double max = score[0];
		
		for (int i = 0; i < classes.length; i++) {
			if (score[i] > max) {
				max = score[i];
				maxi = i;
			}
		}
		
		return maxi;
	}
	
	public static Set<String> extractVocabulary(List<Document> documents) {
		Set<String> vocabulary = new HashSet<String>();
		for (Document d : documents) {
			vocabulary.addAll(d.getTerms());
		}
		return vocabulary;		
	}
	
	public static int countTokens(DocumentsLoader loader, int clazz, String term) {
		List<Document> docs = loader.getTrainingDocuments(clazz);
		int res = 0;		
		for (Document doc : docs) {
			for (String token : doc.getTokens()) {
				if (token.equals(term)) {
					res++;
				}
			}
		}		
		return res;
	}
}